import torch.nn as nn
from transformers import AutoModel
from config import PRETRAINED_DIR


class ReviewAnalyzeModel(nn.Module):
    """ 评论情感分析模型类 """

    def __init__(self, freeze_bert=True):
        super().__init__()
        self.bert = AutoModel.from_pretrained(PRETRAINED_DIR / 'bert-base-chinese')
        self.linear = nn.Linear(in_features=self.bert.config.hidden_size, out_features=1)
        # 是否冻结bert参数
        if freeze_bert:
            for param in self.bert.parameters():
                param.requires_grad = False

    def forward(self, input_ids, attention_mask):
        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        last_hidden_states = outputs.last_hidden_state  # [batch_size, seq_len, hidden_size]
        cls_output = last_hidden_states[:, 0, :]  # [batch_size, hidden_size]
        output = self.linear(cls_output).squeeze(dim=-1)  # [batch_size]
        return output
